#include "server/websocket_session.h"
#include "common/utils.h"
#include <iostream>
#include <cstring>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <netinet/in.h>

namespace websocket {

WebSocketSession::WebSocketSession(int socket) 
    : socket_(socket), connected_(true), closing_(false) {
    
    // 获取客户端地址信息
    sockaddr_in addr{};
    socklen_t len = sizeof(addr);
    if (getpeername(socket, reinterpret_cast<sockaddr*>(&addr), &len) == 0) {
        char ip_str[INET_ADDRSTRLEN];
        inet_ntop(AF_INET, &addr.sin_addr, ip_str, sizeof(ip_str));
        remote_address_ = std::string(ip_str) + ":" + std::to_string(ntohs(addr.sin_port));
    } else {
        remote_address_ = "unknown";
    }
}

WebSocketSession::~WebSocketSession() {
    cleanup();
}

void WebSocketSession::start() {
    read_thread_ = std::thread(&WebSocketSession::readLoop, this);
}

void WebSocketSession::sendText(const std::string& message) {
    if (!connected_ || closing_) return;
    
    WebSocketFrame frame = WebSocketFrame::createTextFrame(message);
    auto data = frame.serialize();
    sendData(data);
}

void WebSocketSession::sendBinary(const std::vector<uint8_t>& data) {
    if (!connected_ || closing_) return;
    
    WebSocketFrame frame(Opcode::BINARY, data);
    auto frame_data = frame.serialize();
    sendData(frame_data);
}

void WebSocketSession::sendPing() {
    if (!connected_ || closing_) return;
    
    WebSocketFrame frame = WebSocketFrame::createPingFrame();
    auto data = frame.serialize();
    sendData(data);
}

void WebSocketSession::close(uint16_t code, const std::string& reason) {
    if (closing_ || !connected_) return;
    
    closing_ = true;
    
    // 发送关闭帧
    WebSocketFrame frame = WebSocketFrame::createCloseFrame(code, reason);
    auto data = frame.serialize();
    sendData(data);
    
    // 关闭socket
    cleanup();
}

void WebSocketSession::readLoop() {
    std::vector<uint8_t> buffer;
    
    while (connected_ && !closing_) {
        // 读取帧头（至少2字节）
        std::vector<uint8_t> header(2);
        if (!readData(header, 2)) {
            break;
        }
        
        // 解析基本帧头
        uint8_t second_byte = header[1];
        bool masked = (second_byte & 0x80) != 0;
        uint64_t payload_len = second_byte & 0x7F;
        
        size_t header_size = 2;
        
        // 处理扩展长度
        std::vector<uint8_t> ext_len_bytes;
        if (payload_len == 126) {
            std::vector<uint8_t> ext_len(2);
            if (!readData(ext_len, 2)) break;
            payload_len = (static_cast<uint64_t>(ext_len[0]) << 8) | static_cast<uint64_t>(ext_len[1]);
            header_size += 2;
            ext_len_bytes = ext_len;
        } else if (payload_len == 127) {
            std::vector<uint8_t> ext_len(8);
            if (!readData(ext_len, 8)) break;
            payload_len = 0;
            for (int i = 0; i < 8; ++i) {
                payload_len = (payload_len << 8) | static_cast<uint64_t>(ext_len[i]);
            }
            header_size += 8;
            ext_len_bytes = ext_len;
        }
        
        // 读取masking key
        std::vector<uint8_t> masking_key;
        if (masked) {
            masking_key.resize(4);
            if (!readData(masking_key, 4)) break;
            header_size += 4;
        }
        
        // 读取payload
        std::vector<uint8_t> payload(payload_len);
        if (payload_len > 0 && !readData(payload, payload_len)) {
            break;
        }
        
        // 构建完整帧数据并解析（保留原始掩码与扩展长度，交由 parse 处理解码）
        std::vector<uint8_t> frame_data;
        frame_data.insert(frame_data.end(), header.begin(), header.end());
        if (!ext_len_bytes.empty()) {
            frame_data.insert(frame_data.end(), ext_len_bytes.begin(), ext_len_bytes.end());
        }
        if (masked && !masking_key.empty()) {
            frame_data.insert(frame_data.end(), masking_key.begin(), masking_key.end());
        }
        frame_data.insert(frame_data.end(), payload.begin(), payload.end());
        
        WebSocketFrame frame;
        if (frame.parse(frame_data)) {
            handleWebSocketFrame(frame);
        }
    }
    
    if (connected_) {
        cleanup();
    }
}

bool WebSocketSession::readData(std::vector<uint8_t>& buffer, size_t size) {
    buffer.resize(size);
    size_t total_read = 0;
    
    while (total_read < size) {
        ssize_t bytes_read = recv(socket_, buffer.data() + total_read, size - total_read, 0);
        if (bytes_read <= 0) {
            return false;
        }
        total_read += bytes_read;
    }
    
    return true;
}

bool WebSocketSession::sendData(const std::vector<uint8_t>& data) {
    size_t total_sent = 0;
    while (total_sent < data.size()) {
        ssize_t bytes_sent = send(socket_, data.data() + total_sent, data.size() - total_sent, 0);
        if (bytes_sent <= 0) {
            return false;
        }
        total_sent += bytes_sent;
    }
    return true;
}

void WebSocketSession::handleWebSocketFrame(const WebSocketFrame& frame) {
    switch (frame.getOpcode()) {
        case Opcode::TEXT:
            if (message_callback_) {
                message_callback_(frame.getPayloadAsString());
            }
            break;
            
        case Opcode::BINARY:
            // 处理二进制数据
            break;
            
        case Opcode::CLOSE:
            close();
            break;
            
        case Opcode::PING:
            // 自动回复PONG
            {
                WebSocketFrame pong_frame = WebSocketFrame::createPongFrame(frame.getPayload());
                auto data = pong_frame.serialize();
                sendData(data);
            }
            break;
            
        case Opcode::PONG:
            // 心跳回复，可以更新最后活动时间
            break;
            
        default:
            break;
    }
}

void WebSocketSession::cleanup() {
    if (connected_) {
        connected_ = false;
        closing_ = true;
        
        if (socket_ >= 0) {
            ::close(socket_);
            socket_ = -1;
        }
        
        if (close_callback_) {
            close_callback_();
        }
    }
    
    // 仅当不在当前读线程内时才 join，避免自我 join 导致崩溃
    if (read_thread_.joinable()) {
        if (std::this_thread::get_id() != read_thread_.get_id()) {
            read_thread_.join();
        } else {
            // 如果在当前线程，避免 join；让线程自然结束
        }
    }
}

} // namespace websocket